import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

def Fun(x,y):
    return x - y + 2*x*x + 2*x*y + y*y

def PxFun(x,y):  # 对 x 求偏导
    return 1 + 4*x + 2*y

def PyFun(x,y):  # 对 y 求偏导
    return -1 + 2*x + 2*y

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')  # 正确初始化 3D 轴

# 生成网格数据
X, Y = np.mgrid[-2:2:40j, -2:2:40j]
Z = Fun(X, Y)

# 绘制 3D 表面
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap="rainbow", alpha=0.7)

# 设置轴标签
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')

# 梯度下降初始化
step = 0.1  # 增大步长
x, y = 0, 0
tag_x, tag_y, tag_z = [x], [y], [Fun(x, y)]
new_x, new_y = x, y

# 梯度下降过程
for _ in range(1000):  # 增加最大迭代次数
    new_x -= step * PxFun(x, y)
    new_y -= step * PyFun(x, y)
    if abs(Fun(x, y) - Fun(new_x, new_y)) < 1e-6:  # 放宽终止条件
        break
    x, y = new_x, new_y
    tag_x.append(x)
    tag_y.append(y)
    tag_z.append(Fun(x, y))

# 绘制梯度下降路径
ax.plot(tag_x, tag_y, tag_z, color='r', linestyle='-', linewidth=2)  # 明确指定颜色和线型

# 标题
plt.title(f"(x,y) ~ ({x:.4f}, {y:.4f})")

# 显示图形
plt.show()
